"""
知识图谱模块 - 处理与知识图谱相关的功能
"""

import streamlit as st
import json
import os
import networkx as nx
import plotly.graph_objects as go

from modules.config import KNOWLEDGE_GRAPH_PATH

@st.cache_data
def load_knowledge_graph():
    """加载知识图谱数据"""
    if os.path.exists(KNOWLEDGE_GRAPH_PATH):
        try:
            with open(KNOWLEDGE_GRAPH_PATH, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            st.error(f"⚠️ 加载知识图谱失败: {str(e)}")
    return None

def get_related_kg_info(artifact_name, kg_data):
    """获取与文物相关的知识图谱信息"""
    if not kg_data:
        return []
    
    try:
        artifact_id = None
        for node in kg_data['nodes']:
            if node['name'] == artifact_name and node['type'] == 'animal':
                artifact_id = node['id']
                break
        
        if not artifact_id:
            return []
        
        relations = []
        
        for edge in kg_data['edges']:
            if edge['source'] == artifact_id:
                target_node = next((node for node in kg_data['nodes'] if node['id'] == edge['target']), None)
                if target_node:
                    relations.append({
                        "relation": edge['type'],
                        "entity": target_node['name'],
                        "type": target_node['type'],
                        "details": edge.get('properties', {})
                    })
            elif edge['target'] == artifact_id:
                source_node = next((node for node in kg_data['nodes'] if node['id'] == edge['source']), None)
                if source_node:
                    relations.append({
                        "relation": edge['type'],
                        "entity": source_node['name'],
                        "type": source_node['type'],
                        "details": edge.get('properties', {})
                    })
        
        return relations
    except Exception as e:
        st.error(f"⚠️ 获取知识图谱信息过程出错: {str(e)}")
        return []

def format_kg_info(relations):
    """格式化知识图谱关系为可读文本"""
    if not relations:
        return ""
    
    try:
        formatted_info = "📚 知识图谱补充信息：\n\n"
        
        type_groups = {}
        for rel in relations:
            entity_type = rel['type']
            if entity_type not in type_groups:
                type_groups[entity_type] = []
            type_groups[entity_type].append(rel)
        
        priority_types = ['habitat', 'food', 'era', 'collection']
        for entity_type in priority_types:
            if entity_type in type_groups:
                relations_of_type = type_groups[entity_type]
                
                type_name = {
                    'animal': '相关动物',
                    'food': '食物',
                    'habitat': '栖息地',
                    'era': '时代',
                    'collection': '动物园'
                }.get(entity_type, entity_type)
                
                formatted_info += f"- **{type_name}**: "
                
                entity_relations = []
                for r in relations_of_type:
                    relation_name = {
                        "lives_in": "栖息于",
                        "eats": "捕食",
                        "exists_in": "存在于哪个时代",
                        "belongs_to": "属于哪个动物园",
                    }.get(r['relation'], r['relation'])
                    
                    entity_relations.append(f"{r['entity']}（{relation_name}）")
                
                formatted_info += "、".join(entity_relations) + "\n"
        
        for entity_type, relations_of_type in type_groups.items():
            if entity_type not in priority_types:
                type_name = {
                    'animal': '相关动物',
                    'food': '食物',
                    'habitat': '栖息地',
                    'era': '时代',
                    'collection': '动物园'
                }.get(entity_type, entity_type)
                
                formatted_info += f"- **{type_name}**: "
                
                entity_relations = []
                for r in relations_of_type:
                    relation_name = {
                        "lives_in": "栖息于",
                        "eats": "捕食",
                        "exists_in": "存在于哪个时代",
                        "belongs_to": "属于哪个动物园"
                    }.get(r['relation'], r['relation'])
                    
                    entity_relations.append(f"{r['entity']}（{relation_name}）")
                
                formatted_info += "、".join(entity_relations) + "\n"
        
        return formatted_info
    except Exception as e:
        st.error(f"⚠️ 格式化知识图谱信息过程出错: {str(e)}")
        return ""

def visualize_artifact_kg(artifact_name, kg_data):
    """可视化与特定文物相关的知识图谱子图"""
    if not kg_data:
        return None
    
    try:
        G = nx.DiGraph()
        
        artifact_node = None
        for node in kg_data['nodes']:
            if node['name'] == artifact_name and node['type'] == 'animal':
                artifact_node = node
                break
        
        if not artifact_node:
            return None
        
        G.add_node(
            artifact_node['id'],
            name=artifact_node['name'],
            type=artifact_node['type'],
            properties=artifact_node.get('properties', {})
        )
        
        related_nodes = set()
        for edge in kg_data['edges']:
            if edge['source'] == artifact_node['id']:
                related_nodes.add(edge['target'])
                G.add_edge(
                    edge['source'],
                    edge['target'],
                    type=edge['type'],
                    properties=edge.get('properties', {})
                )
            elif edge['target'] == artifact_node['id']:
                related_nodes.add(edge['source'])
                G.add_edge(
                    edge['source'],
                    edge['target'],
                    type=edge['type'],
                    properties=edge.get('properties', {})
                )
        
        for node in kg_data['nodes']:
            if node['id'] in related_nodes:
                G.add_node(
                    node['id'],
                    name=node['name'],
                    type=node['type'],
                    properties=node.get('properties', {})
                )
        
        node_types = {
            'animal': {'color': '#8C3D2B', 'label': '动物'},
            'food': {'color': '#308C6C', 'label': '食物'},
            'era': {'color': '#D4A76A', 'label': '时代'},
            'collection': {'color': '#BC6C5A', 'label': '动物园'},
            'habitat': {'color': '#9E6B52', 'label': '栖息地'}
        }
        
        node_x = []
        node_y = []
        node_text = []
        node_color = []
        node_size = []
        node_symbols = []
        
        pos = nx.spring_layout(G, seed=42, k=0.3)
        
        for node_id, coords in pos.items():
            node_x.append(coords[0])
            node_y.append(coords[1])
            
            node_data = G.nodes[node_id]
            node_name = node_data['name']
            node_type = node_data.get('type', 'default')
            
            node_text.append(f"{node_name} ({node_types.get(node_type, {'label': '其他'})['label']})")
            node_color.append(node_types.get(node_type, {'color': '#AAAAAA'})['color'])
            
            if node_id == artifact_node['id']:
                node_size.append(30)
                node_symbols.append('star')
            else:
                node_size.append(20)
                symbol_map = {'animal': 'circle', 'era': 'square', 'collection': 'diamond', 'habitat': 'cross'}
                node_symbols.append(symbol_map.get(node_type, 'circle'))
        
        edge_x = []
        edge_y = []
        edge_text = []
        
        relation_name_map = {
            "lives_in": "栖息于",
            "eats": "捕食",
            "exists_in": "存在于哪个时代",
            "belongs_to": "属于哪个动物园",
        }
        
        for edge in G.edges(data=True):
            x0, y0 = pos[edge[0]]
            x1, y1 = pos[edge[1]]
            
            edge_x.append(x0)
            edge_x.append((x0 + x1) / 2 + (y1 - y0) * 0.1)
            edge_x.append(x1)
            edge_x.append(None)
            
            edge_y.append(y0)
            edge_y.append((y0 + y1) / 2 + (x0 - x1) * 0.1) 
            edge_y.append(y1)
            edge_y.append(None)
            
            relation_type = edge[2]['type']
            relation_label = relation_name_map.get(relation_type, relation_type)
            edge_text.append(relation_label)
        
        edge_trace = go.Scatter(
            x=edge_x, y=edge_y,
            line=dict(width=1, color='rgba(140, 61, 43, 0.5)'),
            hoverinfo='none',
            mode='lines',
            showlegend=False
        )
        
        node_trace = go.Scatter(
            x=node_x, y=node_y,
            mode='markers+text',
            text=node_text,
            textposition="bottom center",
            textfont=dict(family="Source Han Serif CN, Microsoft YaHei, SimHei, sans-serif", size=12),
            marker=dict(
                showscale=False,
                color=node_color,
                size=node_size,
                symbol=node_symbols,
                line=dict(width=1, color='white')
            ),
            hoverinfo='text'
        )
        
        annotations = []
        for i, (x0, y0, x1, y1) in enumerate(zip(
            [pos[edge[0]][0] for edge in G.edges()],
            [pos[edge[0]][1] for edge in G.edges()],
            [pos[edge[1]][0] for edge in G.edges()],
            [pos[edge[1]][1] for edge in G.edges()]
        )):
            if i < len(edge_text):
                annotations.append(dict(
                    x=(x0 + x1) / 2,
                    y=(y0 + y1) / 2,
                    xref="x", yref="y",
                    text=edge_text[i],
                    showarrow=False,
                    font=dict(family="Source Han Serif CN, Microsoft YaHei, SimHei, sans-serif", size=10, color="#333"),
                    bgcolor="rgba(255,255,255,0.8)",
                    borderpad=2,
                    bordercolor="#E0D1B9",
                    borderwidth=1
                ))
        
        fig = go.Figure(data=[edge_trace, node_trace],
                      layout=go.Layout(
                          title={
                              "text": f"《{artifact_name}》相关知识图谱",
                              "font": {"family": "Source Han Serif CN, Microsoft YaHei, SimHei, sans-serif", "size": 18, "color": "#8C3D2B"}
                          },
                          showlegend=False,
                          annotations=annotations,
                          hovermode='closest',
                          margin=dict(b=20, l=5, r=5, t=60),
                          xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                          yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                          plot_bgcolor='white',
                          paper_bgcolor='white'
                      ))
        
        legend_items = []
        for type_name, type_info in node_types.items():
            legend_items.append(
                go.Scatter(
                    x=[None], y=[None],
                    mode='markers',
                    marker=dict(size=10, color=type_info['color']),
                    name=type_info['label'],
                    showlegend=True
                )
            )
        
        for trace in legend_items:
            fig.add_trace(trace)
        
        fig.update_layout(
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=-0.1,
                xanchor="center",
                x=0.5,
                font=dict(family="Source Han Serif CN, Microsoft YaHei, SimHei, sans-serif", size=12)
            )
        )
        
        return fig
    except Exception as e:
        st.error(f"⚠️ 知识图谱可视化过程出错: {str(e)}")
        return None

def create_artifact_info_card(artifact_name, kg_data):
    """创建文物信息卡片"""
    if not kg_data or not artifact_name:
        return None
    
    try:
        artifact_node = None
        for node in kg_data['nodes']:
            if node['name'] == artifact_name and node['type'] == 'animal':
                artifact_node = node
                break
        
        if not artifact_node:
            return None
        
        relations = get_related_kg_info(artifact_name, kg_data)
        
        era = "未知"
        for rel in relations:
            if rel['type'] == 'era' and rel['relation'] in ['exists_in']:
                era = rel['entity']
                break
        
        category = "未知"
        for rel in relations:
            if rel['type'] == 'habitat' and rel['relation'] == 'lives_in':
                category = rel['entity']
                break
        
        collection = "未知"
        for rel in relations:
            if rel['type'] == 'collection' and rel['relation'] == 'belongs_to':
                collection = rel['entity']
                break
        
        creator = "未知"
        for rel in relations:
            if rel['type'] == 'food' and rel['relation'] == 'eats':
                creator = rel['entity']
                break
        
        material = artifact_node.get('properties', {}).get('material', '未知')
        
        info_card = {
            "动物名称": artifact_name,
            "时代": era,
            "栖息地": category,
            "动物园": collection,
            "食物": creator,
            #"材质": material
        }
        
        return info_card
    except Exception as e:
        st.error(f"⚠️ 创建文物信息卡片过程出错: {str(e)}")
        return None 